# model/graphsage_model.py
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import SAGEConv, global_add_pool

class GraphSAGE(nn.Module):
    def __init__(self, in_dim, hidden_dim=128, out_dim=2, num_layers=2, dropout=0.3):
        super(GraphSAGE, self).__init__()
        self.convs = nn.ModuleList()
        self.dropout = dropout

        # 输入层
        self.convs.append(SAGEConv(in_dim, hidden_dim))
        # 隐藏层
        for _ in range(num_layers - 2):
            self.convs.append(SAGEConv(hidden_dim, hidden_dim))
        # 输出层
        self.convs.append(SAGEConv(hidden_dim, hidden_dim))

        # 全连接层
        self.fc = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(p=dropout),
            nn.Linear(hidden_dim // 2, out_dim)
        )

    def forward(self, x, edge_index, batch=None):
        h = x
        for conv in self.convs:
            h = conv(h, edge_index)
            h = F.relu(h)
            h = F.dropout(h, p=self.dropout, training=self.training)

        # 图级池化
        if batch is not None:
            h = global_add_pool(h, batch)
        return self.fc(h)